import numpy as np
import platgo as pg
from scipy.spatial.distance import cdist

"""
------------------------------- Reference --------------------------------
 E. Zitzler, M. Laumanns, and L. Thiele, SPEA2: Improving the strength
 Pareto evolutionary algorithm, Proceedings of the Fifth Conference on
 Evolutionary Methods for Design, Optimization and Control with
 Applications to Industrial Problems, 2001, 95-100.
"""


class SPEA2(pg.Algorithm):

    type: dict = {'single': False, 'multi': True, 'many': False, 'real': True, 'binary': True, 'permutation': True,
                  "large": False, 'expensive': False, 'constrained': False, 'preference': False, 'multimodal': False,
                  'sparse': False, 'gradient': False}

    def __init__(self, maxgen: int, problem: pg.Problem):
        super().__init__(maxgen=maxgen, problem=problem)
        self.name = "SPEA2"
        self.xov = pg.operators.XovSbx()  # 模拟二进制交叉
        self.mut = pg.operators.MutPol(problem)  # 多项式变异

    def go(self, N: int = None, population: pg.Population = None):
        assert N or population, "N and population can't be both None"
        if population is None:
            pop = self.problem.init_pop(N)
        else:
            pop = population
            self.problem.N = pop.decs.shape[0]
        self.problem.cal_obj(pop)
        fitness = cal_fitness(pop)

        while self.not_terminal(pop):
            matingpool = pg.utils.tournament_selection(2, pop.N, fitness)
            p1, p2 = pg.utils.random_selection(pop[matingpool])
            offspring = self.xov(pop, p1, p2)
            offspring = self.mut(offspring)
            self.problem.cal_obj(offspring)  # 计算子代种群目标函数值
            temp_pop = offspring + pop  # 合并种群
            pop, fitness = environmental_selection(temp_pop, pop.N)

        return pop


def cal_fitness(pop: pg.Population) -> np.ndarray:
    N, M = pop.decs.shape
    # 记录支配关系
    dominate = np.zeros((N, N), dtype=bool)
    for i in range(0, N - 1):
        for j in range(i + 1, N):
            k = int(np.any(pop.objv[i] < pop.objv[j])) - int(np.any(pop.objv[i] > pop.objv[j]))
            if k == 1:
                dominate[i][j] = True
            elif k == -1:
                dominate[j][i] = True

    # 计算S(i)
    s = np.sum(dominate, axis=1)

    # 计算R(i)
    r = np.zeros(N)
    for i in range(N):
        r[i] = np.sum(s[dominate[:, i]])

    # 计算D(i)
    distance = cdist(pop.objv, pop.objv)
    distance[np.eye(len(distance), dtype=bool)] = np.inf
    distance = np.sort(distance, axis=1)
    d = 1 / (distance[:, int(np.floor(np.sqrt(N)) - 1)] + 2)
    # 计算fitness
    fitness = r + d

    return fitness


def truncation(objv: np.ndarray, K: int):
    # 截断策略
    distance = cdist(objv, objv)
    distance[np.eye(len(distance), dtype=bool)] = np.inf
    delete = np.zeros(len(objv), dtype=bool)

    while np.sum(delete) < K:
        remain = np.argwhere(~delete).flatten()
        temp = distance[remain]
        temp = temp[:, remain]
        temp = np.sort(temp, axis=1)
        _, rank = np.unique(temp, return_index=True, axis=0)
        delete[remain[rank[0]]] = True

    return delete


def environmental_selection(pop: pg.Population, N: int) -> tuple:
    fitness = cal_fitness(pop)
    next = fitness < 1
    if np.sum(next) < N:
        rank = np.argsort(fitness)
        next[rank[:N]] = True
    elif np.sum(next) > N:
        delete = truncation(pop.objv[next], np.sum(next) - N)
        temp = np.argwhere(next)
        next[temp[delete]] = False

    pop = pop[next]
    fitness = fitness[next]

    return pop, fitness
